from typing import Any, List, Optional

import requests
from langchain_core.embeddings import Embeddings
from langchain_core.utils import (
    secret_from_env,
)
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    SecretStr,
    model_validator,
)
from requests import RequestException
from typing_extensions import Self
from patagent.constant import PATSNAP_API_URL
import base64
import struct


class PatsnapEmbeddings(BaseModel, Embeddings):
    """Patsnap Embedding models.

    Setup:
        To use, you should set the environment variable ``PATSNAP_API_KEY`` to
        your API key or pass it as a named parameter to the constructor.

        .. code-block:: bash

            export PATSNAP_API_KEY="your-api-key"

    Instantiate:
        .. code-block:: python

            from embedding import PatsnapEmbeddings

            embeddings = PatsnapEmbeddings()

    Embed:
        .. code-block:: python

            # embed the documents
            vectors = embeddings.embed_documents([text1, text2, ...])

            # embed the query
            vectors = embeddings.embed_query(text)
    """  # noqa: E501

    session: Any = None  #: :meta private:
    model_name: str = Field(default="", alias="model")
    """The model used to embed the documents."""
    patsnap_api_key: SecretStr = Field(
        alias="api_key",
        default_factory=secret_from_env(["PATSNAP_API_KEY"]),
    )
    rt_768: int = 1
    model_config = ConfigDict(populate_by_name=True, protected_namespaces=())

    @model_validator(mode="after")
    def validate_environment(self) -> Self:
        """Validate that auth token exists in environment."""
        session = requests.Session()
        session.headers.update(
            {
                "Authorization": f"Bearer {self.patsnap_api_key.get_secret_value()}",
                "X-PatSnap-Version": "v1",
                "Content-type": "application/json",
            }
        )
        self.session = session
        return self

    def _embed(self, texts: List[str]) -> Optional[List[List[float]]]:
        """Internal method to call Patsnap Embedding API and return embeddings.

        Args:
            texts: A list of texts to embed.

        Returns:
            A list of list of floats representing the embeddings, or None if an
            error occurs.
        """
        embed_results = []
        for text in texts:
            response = self.session.post(
                PATSNAP_API_URL, json={'data': {"text": text, "rt_768": self.rt_768}}
            )
            # Raise exception if response status code from 400 to 600
            response.raise_for_status()
            # Check if the response status code indicates success
            if response.status_code == 200:
                resp = response.json()
                vector = resp['data']['vec_768']
                embeddings = struct.unpack("768f", base64.b64decode(vector))
                # Return just the embeddings
                embed_results.append(list(embeddings))
            else:
                # Log error or handle unsuccessful response appropriately
                # Handle 100 <= status_code < 400, not include 200
                raise RequestException(
                    f"Error: Received status code {response.status_code} from "
                    "`PatsnapEmbedding` API"
                )
        return embed_results

    def embed_documents(self, texts: List[str]) -> Optional[List[List[float]]]:  # type: ignore[override]
        """Public method to get embeddings for a list of documents.

        Args:
            texts: The list of texts to embed.

        Returns:
            A list of embeddings, one for each text, or None if an error occurs.
        """
        return self._embed(texts)

    def embed_query(self, text: str) -> Optional[List[float]]:  # type: ignore[override]
        """Public method to get embedding for a single query text.

        Args:
            text: The text to embed.

        Returns:
            Embeddings for the text, or None if an error occurs.
        """
        result = self._embed([text])
        return result[0] if result is not None else None
